#!/usr/bin/env python
# coding: utf-8

# In[ ]:


######数据预处理######
###加载与清洗数据

import pandas as pd
import numpy as np
from biom import load_table
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest, f_classif, VarianceThreshold
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.pipeline import Pipeline
import matplotlib.pyplot as plt
import seaborn as sns


# 加载宏基因组物种数据（BIOM格式）和宏转录组数据
def load_biom_to_df(biom_path):
    biom_table = load_table(biom_path)
    return pd.DataFrame(
        biom_table.matrix_data.toarray().T,
        index=biom_table.ids(axis='sample'),
        columns=biom_table.ids(axis='observation')
    )


df_pathway = pd.read_csv("pathabundance_relab_1.tsv", sep="\t", index_col=0).T
df_ec = pd.read_csv("ecs_relab.tsv", sep="\t", index_col=0).T
df_pathway_rna = pd.read_csv("pathabundance_relab_2.tsv", sep="\t", index_col=0).T
df_taxa = load_biom_to_df("taxonomic_profiles.biom")

print("=== 正确数据结构应如下 ===")
print("Taxa (样本×物种):", df_taxa.shape)
print("Pathway (样本×通路):", df_pathway.shape)
print("EC (样本×酶):", df_ec.shape)
print("Pathway RNA (样本×通路):", df_pathway_rna.shape)



# 获取共有的样本ID
common_samples = set(df_taxa.index) & set(df_pathway.index) & set(df_ec.index) & set(df_pathway_rna.index)

# 筛选数据
dfs = []
for df in [df_taxa, df_pathway, df_ec, df_pathway_rna]:
    dfs.append(df.loc[list(common_samples)])  # 转换为list避免警告

df_taxa, df_pathway, df_ec, df_pathway_rna = dfs

print(f"成功对齐{len(common_samples)}个样本")




np.random.seed(42)
y = pd.Series(np.random.randint(0, 2, size=len(common_samples)), 
             index=common_samples)  # 二分类标签

##### 单组学分析 ######
# 定义单组学数据集
single_omics = {
    "Taxonomy": df_taxa,
    "Pathway": df_pathway,
    "EC": df_ec,
    "Pathway_RNA": df_pathway_rna
}

cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
# 单组学评估函数
def evaluate_single_omics(X, y, model):
    pipe = Pipeline([
        ('scaler', StandardScaler()),
        ('variance_threshold', VarianceThreshold(threshold=0.01)),  # 添加方差过滤
        ('selector', SelectKBest(f_classif, k=100)),  # 保持与多组学相同的特征选择
        ('model', model)
    ])
    scores = cross_val_score(
        pipe, X, y, cv=cv, 
        scoring='roc_auc',
        n_jobs=-1
    )
    return {
        'AUC_mean': np.mean(scores),
        'AUC_std': np.std(scores),
        'model': pipe.fit(X, y)  # 保存训练好的模型
    }

# 评估所有单组学
single_results = {}
for name, data in single_omics.items():
    print(f"\n正在评估单组学: {name}...")
    single_results[name] = evaluate_single_omics(
        data.loc[common_samples], 
        y,
        RandomForestClassifier(n_estimators=200, max_depth=10, random_state=42)
    )
    print(f"{name} AUC: {single_results[name]['AUC_mean']:.3f} ± {single_results[name]['AUC_std']:.3f}")



    

###### 多组学分析 ###### 


###数据标准化与特征选择
# 合并多组学数据
X = pd.concat([
    df_taxa.add_prefix("taxa_"),
    df_pathway.add_prefix("pathway_"),
    df_ec.add_prefix("ec_"),
    df_pathway_rna.add_prefix("pathway_rna_")
], axis=1)

# 标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# 特征选择 
# 选择前100个最重要的特征
selector = SelectKBest(score_func=f_classif, k=100)
X_selected = selector.fit_transform(X_scaled, y)

# 获取选中的特征名
selected_features = X.columns[selector.get_support()]
print("选中的Top 10特征:", selected_features[:10])

#数据验证
print("\n最终数据形状:")
print("X_selected:", X_selected.shape)  
print("y:", y.shape)  




######多组学整合建模######
###早期整合（特征拼接+随机森林）
# 初始化模型和交叉验证
model_rf = RandomForestClassifier(
    n_estimators=200,
    max_depth=10,
    random_state=42
)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

# 交叉验证评估
scores = cross_val_score(
    model_rf, 
    X_selected, 
    y, 
    cv=cv, 
    scoring='roc_auc',
    n_jobs=-1
)
print(f"Early Integration (RF) AUC: {np.mean(scores):.3f} ± {np.std(scores):.3f}")

# 训练最终模型
final_model = model_rf.fit(X_selected, y)


multi_results = {
    'Multi-omics': {
        'AUC_mean': np.mean(scores),
        'AUC_std': np.std(scores),
        'model': final_model
    }
}

###### 结果比较与可视化 ######
# 合并结果
all_results = {**single_results, **multi_results}
results_df = pd.DataFrame({
    'Dataset': all_results.keys(),
    'AUC': [x['AUC_mean'] for x in all_results.values()],
    'AUC_std': [x['AUC_std'] for x in all_results.values()]
})
# 性能对比条形图
plt.figure(figsize=(10, 6))
sns.barplot(
    data=results_df.sort_values('AUC', ascending=False),
    x='Dataset', y='AUC',
    yerr=results_df['AUC_std'],
    palette='viridis'
)
plt.ylim(0.5, max(results_df['AUC']) + 0.1)
plt.title("Single-omics vs Multi-omics Performance (AUC)")
plt.axhline(y=0.5, color='red', linestyle='--', label='Random Guess')
plt.legend()
plt.tight_layout()
plt.savefig('performance_comparison.png', dpi=300)
plt.show()

